import cv2
from ultralytics import YOLO

# 加载 YOLO 模型
model = YOLO("../yolo11x_pt/yolo11n.pt")  # 加载预训练模型

# 打开视频流（可以是摄像头或视频文件）
cap = cv2.VideoCapture(0)  # 0 表示摄像头，或替换为视频文件路径

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 使用 YOLO 模型进行检测
    results = model.predict(source=frame, save=False, conf=0.5)

    # 获取检测结果
    for result in results[0].boxes:
        x1, y1, x2, y2 = map(int, result.xyxy[0])  # 边界框坐标
        conf = result.conf[0]  # 置信度
        cls = int(result.cls[0])  # 类别索引
        label = f"{model.names[cls]} {conf:.2f}"  # 类别名称和置信度

        # 绘制边界框和标签
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2
        )

    # 显示检测结果
    cv2.imshow("YOLO Detection", frame)

    # 按 'q' 键退出
    if cv2.waitKey(1) & 0xFF == ord("q"):
        break

cap.release()
cv2.destroyAllWindows()
